/** @file admm_result.h
 *
 *  @author Dongryeol Lee (dongryel@cc.gatech.edu)
 */

#ifndef MLPACK_ADMM_ADMM_RESULT_H
#define MLPACK_ADMM_ADMM_RESULT_H

#include <armadillo>
#include <boost/bind.hpp>
#include <boost/math/distributions/normal.hpp>
#include <boost/mpi.hpp>
#include <boost/serialization/serialization.hpp>
#include "mlpack/admm/admm_arguments.h"

namespace mlpack {
namespace admm {

/** @brief Represents the storage of ADMM computation results.
 */
class ADMMResult {
  private:

    // For BOOST serialization.
    friend class boost::serialization::access;

  public:

    arma::vec local_adjusted_lagrangian_variables_;

    arma::vec local_dual_variables_;

    arma::vec primal_variables_;

    arma::vec synchronized_lagrangian_variables_;

    arma::vec query_predictions_;

  public:

    template<typename ADMMTableType>
    void Predict(const ADMMTableType &query_table) {
      query_predictions_.zeros(query_table.n_entries());
      for(int i = 0; i < query_table.n_entries(); i++) {

        // Hard-coded for SVM at this point.
        arma::vec query_point;
        arma::vec query_weight;
        query_table.get(i, &query_point, &query_weight);
        query_predictions_[i] =
          (arma::dot(primal_variables_, query_point) > 0) ? 1.0 : -1.0;
      }
    }

    void Export(
      const std::string & model_output_file,
      const std::string * prediction_output_file) {
      FILE *foutput = fopen(model_output_file.c_str(), "w+");
      for(int i = 0; i < static_cast<int>(primal_variables_.n_elem); i++) {
        fprintf(foutput, "%g\n", primal_variables_[i]);
      }
      std::cerr << "Outputted the model to " << model_output_file << "...\n";
      fclose(foutput);
      if(prediction_output_file != NULL) {
        foutput = fopen(prediction_output_file->c_str(), "w+");
        for(int i = 0; i < static_cast<int>(query_predictions_.n_elem); i++) {
          fprintf(foutput, "%g\n", query_predictions_[i]);
        }
        fclose(foutput);
      }
    }

    void Init(int dimension_in, int num_points_in) {
      local_adjusted_lagrangian_variables_.zeros(dimension_in);
      local_dual_variables_.zeros(num_points_in);
      primal_variables_.zeros(dimension_in);
      synchronized_lagrangian_variables_.zeros(dimension_in);
    }
};
}
}

#endif
